{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "from chromadb.utils import embedding_functions\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "# 初始化 Chroma 客户端\n",
    "chroma_client = chromadb.Client()\n",
    "\n",
    "# 创建或加载一个集合（collection）\n",
    "collection = chroma_client.get_or_create_collection(name=\"documents\")\n",
    "\n",
    "# 初始化嵌入模型\n",
    "embedding_model = SentenceTransformer('all-MiniLM-L6-v2')\n",
    "\n",
    "def add_document_to_db(document_text, document_id):\n",
    "    # 生成文档的嵌入向量\n",
    "    embedding = embedding_model.encode(document_text)\n",
    "    # 将文档添加到 Chroma 集合\n",
    "    collection.add(\n",
    "        documents=[document_text],\n",
    "        embeddings=[embedding],\n",
    "        ids=[document_id]\n",
    "    )\n",
    "\n",
    "def retrieve_documents(query, top_k=3):\n",
    "    # 生成查询的嵌入向量\n",
    "    query_embedding = embedding_model.encode(query)\n",
    "    # 检索最相关的文档\n",
    "    results = collection.query(\n",
    "        query_embeddings=[query_embedding],\n",
    "        n_results=top_k\n",
    "    )\n",
    "    return results[\"documents\"][0]  # 返回最相关的文档"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ollama\n",
    "\n",
    "def generate_qa(query):\n",
    "    # 从 Chroma 中检索相关文档\n",
    "    relevant_docs = retrieve_documents(query)\n",
    "    context = \"\\n\".join(relevant_docs)  # 将检索到的文档拼接为上下文\n",
    "\n",
    "    # 使用 Ollama 生成问题和答案\n",
    "    response = ollama.generate(\n",
    "        model=\"your-model-name\",  # 替换为你的模型名称\n",
    "        prompt=f\"根据以下上下文生成问题和答案：\\n{context}\\n\\n问题：{query}\"\n",
    "    )\n",
    "    return response[\"output\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "\n",
    "def process_document(file):\n",
    "    with open(file.name, \"r\", encoding=\"utf-8\") as f:\n",
    "        document_text = f.read()\n",
    "    # 将文档添加到 Chroma 数据库\n",
    "    document_id = f\"doc_{hash(document_text)}\"  # 生成唯一 ID\n",
    "    add_document_to_db(document_text, document_id)\n",
    "    return \"文档已成功上传并添加到知识库！\"\n",
    "\n",
    "def generate_qa_interface(query):\n",
    "    qa_pairs = generate_qa(query)\n",
    "    return qa_pairs\n",
    "\n",
    "def evaluate_answers(user_answers, correct_answers):\n",
    "    # 简单的答案评判逻辑\n",
    "    correct_count = 0\n",
    "    for user_answer, correct_answer in zip(user_answers, correct_answers):\n",
    "        if user_answer.strip().lower() == correct_answer.strip().lower():\n",
    "            correct_count += 1\n",
    "    return f\"你答对了 {correct_count} 道题，共 {len(correct_answers)} 道题。\"\n",
    "\n",
    "with gr.Blocks() as demo:\n",
    "    with gr.Tab(\"知识库\"):\n",
    "        file_input = gr.File(label=\"上传文档\")\n",
    "        upload_status = gr.Textbox(label=\"上传状态\")\n",
    "        upload_button = gr.Button(\"上传文档\")\n",
    "        upload_button.click(process_document, inputs=file_input, outputs=upload_status)\n",
    "\n",
    "        query_input = gr.Textbox(label=\"输入问题\")\n",
    "        qa_output = gr.Textbox(label=\"生成的问题和答案\", lines=10)\n",
    "        generate_button = gr.Button(\"生成问题和答案\")\n",
    "        generate_button.click(generate_qa_interface, inputs=query_input, outputs=qa_output)\n",
    "\n",
    "    with gr.Tab(\"考试评判\"):\n",
    "        user_answers = gr.Textbox(label=\"你的答案（每行一个答案）\", lines=10)\n",
    "        correct_answers = gr.Textbox(label=\"正确答案（每行一个答案）\", lines=10)\n",
    "        evaluate_button = gr.Button(\"评判答案\")\n",
    "        result_output = gr.Textbox(label=\"评判结果\")\n",
    "        evaluate_button.click(evaluate_answers, inputs=[user_answers, correct_answers], outputs=result_output)\n",
    "\n",
    "demo.launch()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
